#!/usr/bin/env python

import cPickle
import os
import optparse
import random
import sys

import numpy
import tables

import charmicat
import histogram


def main():
    '''
    python -m batchcat.run.module_hists module_name paths.pkl hist_names.pkl

    Batch-processing module for `bcat_module_hists.py'.
    '''

    op = optparse.OptionParser(usage=main.__doc__)

    opts, args = op.parse_args()

    try:
        module_name = args.pop(0)
    except IndexError:
        msg = "Please supply the analysis module name, "
        msg += "a pickle file containing the input paths, "
        msg += "and a pickle file containing the histogram names to be filled."
        raise charmicat.BadCommandLineError(msg, -1)

    try:
        paths_file = args.pop(0)
    except IndexError:
        msg = "Please supply "
        msg += "a pickle file containing the input paths "
        msg += "and a pickle file containing the histogram names to be filled."
        raise charmicat.BadCommandLineError(msg, -2)

    try:
        hnames_file = args.pop(0)
    except IndexError:
        msg = "Please supply "
        msg += "a pickle file containing the histogram names to be filled."
        raise charmicat.BadCommandLineError(msg, -3)
    
    try:
        module = __import__(module_name)           
    except ImportError, e:
        msg = "Error importing module `{0}': {1}".format(module_name, e)
        raise charmicat.BadCommandLineError(msg, -4)
    
    with open(paths_file) as fp:
        paths = cPickle.load(fp)

    with open(hnames_file) as fp:
        hcol_names, hvec_names = cPickle.load(fp)

    
    hcol = {}
    hcol['1d'] = [module.hists[hn] for hn in hcol_names if len(module.hists[hn].assoc_cols) == 1]
    hcol['2d'] = [module.hists[hn] for hn in hcol_names if len(module.hists[hn].assoc_cols) == 2]
    hcol['3d'] = [module.hists[hn] for hn in hcol_names if len(module.hists[hn].assoc_cols) == 3]

    hvec = [module.hists[hn] for hn in hvec_names]

    if len(hcol_names) != 0:
        bufsize = max([module.hists[hn]._bufsize for hn in hcol_names])

    for ip, path in enumerate(paths):
        fpin = tables.openFile(path)

        # Begin column-based filling
        if len(hcol_names) != 0:
            nrows = len(fpin.root.nt)
            for i in xrange(0, nrows, bufsize):
                stop = i + bufsize
                if (i + bufsize) > nrows:
                    stop = nrows

                rows = fpin.root.nt.read(start=i, stop=stop)

                assoc_cols = charmicat.histogram_assoc_cols(module, hcol_names)

                for f in assoc_cols:
                    rows[:][f] = numpy.nan_to_num(rows[:][f])                            

                for h in hcol['1d']:
                    if h.name in module.hist_colfilters:
                        h.fill(module.hist_colfilters[h.name](rows[:][h.assoc_cols[0]]))
                    else:
                        h.fill(rows[:][h.assoc_cols[0]])

                for h in hcol['2d']:
                    h.fill(zip(rows[:][h.assoc_cols[0]], 
                               rows[:][h.assoc_cols[1]]))

                for h in hcol['3d']:
                    if h.name in module.hists.recompute:
                        h.fill(module.hists.recompute[h.name](rows))
                    else:
                        h.fill(zip(rows[:][h.assoc_cols[0]], 
                                   rows[:][h.assoc_cols[1]], 
                                   rows[:][h.assoc_cols[2]]))

        # Begin four-vector-based filling
        if len(hvec_names) != 0:
            vecs = module.fill.four_vectors(fpin.root.nt)

            for h in hvec:
                h.fill(module.fill.__dict__[h.name](fpin.root.nt, vecs))

        fpin.close()

    hout = {}
    hout.update(dict([(h.name, h) for h in hcol['1d']]))
    hout.update(dict([(h.name, h) for h in hcol['2d']]))
    hout.update(dict([(h.name, h) for h in hcol['3d']]))
    hout.update(dict([(h.name, h) for h in hvec]))

    outpath = os.path.join(os.path.dirname(paths_file), 
                           os.path.basename(paths_file).rsplit('_', 1)[0]) + '.hpt'

    histogram.dump(hout, outpath)


if __name__ == "__main__":
    try:
        sys.exit(main())
    except charmicat.BadCommandLineError, e:
        import os
        sys.stderr.write(str(e) + "\n")
        sys.stderr.write("See `{0} -h' for more.\n".format(os.path.basename(sys.argv[0])))
        sys.stderr.flush()
        sys.exit(e.retcode)
